import torch
import torch.nn as nn
import tqdm
from torch.optim import Adam
from torch.utils.data import DataLoader

from BERT.bert_pytorch.model import BERTLM, BERT
from BERT.bert_pytorch.trainer.optim_schedule import ScheduledOptim


class BERTTrainer:
    """
    BERTTrainer 使用两种 LM 训练方法预训练 BERT 模型。

        1. Masked Language Model (掩码语言模型): 3.3.1 Task #1: Masked LM
        2. Next Sentence prediction (下一句预测): 3.3.2 Task #2: Next Sentence Prediction

    请查看 README.md 了解详细信息和简单示例。

    """

    def __init__(self, bert: BERT, vocab_size: int,
                 train_dataloader: DataLoader, test_dataloader: DataLoader = None,
                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000,
                 with_cuda: bool = True, cuda_devices=None, log_freq: int = 10):
        """
        :param bert: 你想要训练的 BERT 模型
        :param vocab_size: 总词汇表大小
        :param train_dataloader: 训练数据集数据加载器
        :param test_dataloader: 测试数据集数据加载器 [可以为 None]
        :param lr: 优化器的学习率
        :param betas: Adam 优化器的 betas 参数
        :param weight_decay: Adam 优化器的权重衰减参数
        :param with_cuda: 使用 CUDA 进行训练
        :param log_freq: 批次迭代的日志记录频率
        """

        # 为 BERT 训练设置 cuda 设备，参数 -c, --cuda 应该为 true
        cuda_condition = torch.cuda.is_available() and with_cuda
        self.device = torch.device("cuda:0" if cuda_condition else "cpu")

        # 这个 BERT 模型将每 epoch 保存一次
        self.bert = bert
        # 初始化 BERT 语言模型，使用 BERT 模型
        self.model = BERTLM(bert, vocab_size).to(self.device)

        # 如果 CUDA 检测到超过 1 个 GPU，则进行分布式 GPU 训练
        if with_cuda and torch.cuda.device_count() > 1:
            print("Using %d GPUS for BERT" % torch.cuda.device_count())
            self.model = nn.DataParallel(self.model, device_ids=cuda_devices)

        # 设置训练和测试数据加载器
        self.train_data = train_dataloader
        self.test_data = test_dataloader

        # 使用超参数设置 Adam 优化器
        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
        self.optim_schedule = ScheduledOptim(self.optim, self.bert.hidden, n_warmup_steps=warmup_steps)

        # 使用负对数似然损失函数来预测 masked_token
        self.criterion = nn.NLLLoss(ignore_index=0)

        self.log_freq = log_freq

        print("Total Parameters:", sum([p.nelement() for p in self.model.parameters()]))

    def train(self, epoch):
        self.iteration(epoch, self.train_data)

    def test(self, epoch):
        self.iteration(epoch, self.test_data, train=False)

    def iteration(self, epoch, data_loader, train=True):
        """
        循环遍历 data_loader 进行训练或测试
        如果在训练状态，则激活反向传播操作
        并且还会每 epoch 自动保存模型

        :param epoch: 当前 epoch 索引
        :param data_loader: torch.utils.data.DataLoader 用于迭代
        :param train: 是否为训练或测试的布尔值
        :return: None
        """
        str_code = "train" if train else "test"

        # 设置 tqdm 进度条
        data_iter = tqdm.tqdm(enumerate(data_loader),
                              desc="EP_%s:%d" % (str_code, epoch),
                              total=len(data_loader),
                              bar_format="{l_bar}{r_bar}")

        avg_loss = 0.0
        total_correct = 0
        total_element = 0

        for i, data in data_iter:
            # 0. batch_data 将被发送到设备（GPU 或 cpu）
            data = {key: value.to(self.device) for key, value in data.items()}

            # 1. 前向传播 next_sentence_prediction 和 masked_lm 模型
            next_sent_output, mask_lm_output = self.model.forward(data["bert_input"], data["segment_label"])

            # 2-1. is_next 分类结果的 NLL（负对数似然）损失
            next_loss = self.criterion(next_sent_output, data["is_next"])

            # 2-2. 预测 masked token 词的 NLLLoss
            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data["bert_label"])

            # 2-3. 添加 next_loss 和 mask_loss : 3.4 预训练过程
            loss = next_loss + mask_loss

            # 3. 仅在训练时进行反向传播和优化
            if train:
                self.optim_schedule.zero_grad()
                loss.backward()
                self.optim_schedule.step_and_update_lr()

            # 下一句预测准确率
            correct = next_sent_output.argmax(dim=-1).eq(data["is_next"]).sum().item()
            avg_loss += loss.item()
            total_correct += correct
            total_element += data["is_next"].nelement()

            post_fix = {
                "epoch": epoch,
                "iter": i,
                "avg_loss": avg_loss / (i + 1),
                "avg_acc": total_correct / total_element * 100,
                "loss": loss.item()
            }

            if i % self.log_freq == 0:
                data_iter.write(str(post_fix))

        print("EP%d_%s, avg_loss=" % (epoch, str_code), avg_loss / len(data_iter), "total_acc=",
              total_correct * 100.0 / total_element)

    def save(self, epoch, file_path="output/bert_trained.model"):
        """
        在 file_path 上保存当前的 BERT 模型

        :param epoch: 当前 epoch 编号
        :param file_path: 模型输出路径，将变为 file_path+"ep%d" % epoch
        :return: final_output_path
        """
        output_path = file_path + ".ep%d" % epoch
        torch.save(self.bert.state_dict(), output_path)  # 修改为保存 state_dict
        self.bert.to(self.device)  # 可选: 如果你希望保存后模型仍然在 GPU 上，保留这行，否则可以删除
        print("EP:%d Model state_dict Saved on:" % epoch, output_path)  # 修改打印信息
        return output_path
